# get_network_ip/ping_utils/ping_utils.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import re
import yaml
import time
import argparse
import subprocess
from datetime import datetime
from typing import Iterable, Dict, List
from concurrent.futures import ThreadPoolExecutor, as_completed

# 优先：安装后的 share 路径（ros2 pkg prefix 可找到）
def _installed_devices_yaml() -> str:
    try:
        from ament_index_python.packages import get_package_share_directory
        pkg_share = get_package_share_directory('get_network_ip')
        return os.path.join(pkg_share, 'config', 'devices.yaml')
    except Exception:
        return ""

# 兜底：源码树相对路径
def _source_devices_yaml() -> str:
    # ping_utils.py 在 get_network_ip/ping_utils/ 下，回退到 ../config/devices.yaml
    return os.path.normpath(os.path.join(os.path.dirname(__file__), '..', 'config', 'devices.yaml'))

# 默认 YAML 路径（自动选择）
def _default_devices_yaml() -> str:
    p = _installed_devices_yaml()
    return p if p and os.path.isfile(p) else _source_devices_yaml()

# ----------------- 核心 ping -----------------
def _ping_once(ip: str, timeout: float = 1.0) -> dict:
    try:
        res = subprocess.run(
            ["ping", "-c", "1", "-W", str(int(timeout)), ip],
            capture_output=True, text=True
        )
        if res.returncode != 0:
            return {"online": False, "rtt_ms": None}
        m = re.search(r'time[=<]\s*([\d.]+)\s*ms', res.stdout)
        rtt = float(m.group(1)) if m else None
        return {"online": True, "rtt_ms": rtt}
    except Exception:
        return {"online": False, "rtt_ms": None}

def ping_hosts(ips: Iterable[str], timeout: float = 1.0, max_workers: int = 20) -> Dict[str, dict]:
    ip_list = list(dict.fromkeys(str(ip).strip() for ip in ips if str(ip).strip()))
    if not ip_list:
        return {}
    results: Dict[str, dict] = {}
    with ThreadPoolExecutor(max_workers=max_workers) as ex:
        fmap = {ex.submit(_ping_once, ip, timeout): ip for ip in ip_list}
        for fut in as_completed(fmap):
            ip = fmap[fut]
            try:
                results[ip] = fut.result()
            except Exception:
                results[ip] = {"online": False, "rtt_ms": None}
    return results

# ----------------- YAML 读取与组装 -----------------
def load_devices_from_yaml(path: str | None = None) -> dict:
    """读取 YAML，返回 period_sec/timeout/topic/devices。path 为空则用默认。"""
    path = path or _default_devices_yaml()
    if not os.path.isfile(path):
        raise FileNotFoundError(f"配置文件不存在: {path}")

    with open(path, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f) or {}

    period_sec = float(cfg.get("period_sec", 1.0))
    timeout = float(cfg.get("timeout", 1.0))
    topic_name = cfg.get("topic_name", "/network/devices_status")

    devices = []
    for d in cfg.get("devices", []):
        name = str(d.get("name", "")).strip()
        ip = str(d.get("ip", "")).strip()
        if name and ip:
            devices.append({"name": name, "ip": ip})

    return {
        "period_sec": period_sec,
        "timeout": timeout,
        "topic_name": topic_name,
        "devices": devices,
        "config_path": path,
    }

def ping_from_yaml(path: str | None = None, max_workers: int = 20) -> List[dict]:
    """读 YAML → 并发检测 → 返回可直接发布的 payload 列表。"""
    cfg = load_devices_from_yaml(path)
    ips = [d["ip"] for d in cfg["devices"]]
    result_map = ping_hosts(ips, timeout=cfg["timeout"], max_workers=max_workers)

    now = datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ")
    payload: List[dict] = []
    for d in cfg["devices"]:
        info = result_map.get(d["ip"], {"online": False, "rtt_ms": None})
        payload.append({
            "name": d["name"],
            "ip": d["ip"],
            "online": bool(info.get("online")),
            "rtt_ms": info.get("rtt_ms"),
            "last_updated": now
        })
    return payload

# ----------------- 测试入口（独立运行） -----------------
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="并发 Ping 工具（独立运行测试）")
    parser.add_argument("--yaml", help="devices.yaml 路径；缺省则自动选择默认路径")
    parser.add_argument("--interval", type=float, default=1.0, help="刷新间隔秒")
    parser.add_argument("--timeout", type=float, default=None, help="覆盖 YAML 的 timeout（秒）")
    args = parser.parse_args()

    cfg = load_devices_from_yaml(args.yaml)
    if args.timeout is not None:
        cfg["timeout"] = float(args.timeout)

    print(f"开始持续网络检测（Ctrl+C 退出）... 使用配置: {cfg['config_path']}\n")
    try:
        while True:
            ips = [d["ip"] for d in cfg["devices"]]
            results = ping_hosts(ips, timeout=cfg["timeout"])
            now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            print(f"\n[{now}] 检测结果（{len(ips)} 台设备）：")
            for d in cfg["devices"]:
                info = results.get(d["ip"], {"online": False, "rtt_ms": None})
                online = "✅ 在线" if info["online"] else "❌ 离线"
                rtt = f"{info['rtt_ms']:.2f} ms" if info["rtt_ms"] is not None else "-"
                print(f"  {d['name']:<12} {d['ip']:<15} {online:6} RTT: {rtt}")
            time.sleep(args.interval)
    except KeyboardInterrupt:
        print("\n已停止检测。")
